import time

import numpy as np
import cv2 as cv
import os
import os.path
import cv2
from utils.MQ.message import send_message
import json


class Preprocessing():
    def __init__(self, MQ_send, connection_send, task_id, task_type, npy_file, png_file, area, targetFolder='./out'):
        """
        class to preprocess meta png and npy
        Auth: WZW
        Args:
            self:
            MQ_send:
            connection_send:
            task_id:
            task_type:
            npy_file:  file obsolete path
            png_file:  file obsolete path
            area: side is 0，up is 1
            targetFolder: folder to save generated png

        Returns:

        """
        self.MQ_send = MQ_send
        self.connection_send = connection_send
        self.task_id = task_id
        self.task_type = task_type

        self.img_light_pth = png_file
        self.img_light = cv2.imread(png_file)
        self.npy_height_pth = npy_file
        self.npy_height = np.load(npy_file)
        self.area = area  # 0是侧面，1是正面
        self.target_folder = targetFolder
        self.img_16bit = np.array([0])
        self.img_gray8 = np.array([0])
        self.img_max = np.array([0])
        self.img_min = np.array([0])
        self.img_rgb = np.array([0])

    def clean_light(self):
        """
            亮度图清洗，去除黑色部分，只保留工件
            必须传入包括文件名的文件路径
            Args:
                img_name: image name with file path
                imshow:
                imsave:
                target_folder: folder to save image

            Returns:
                rgb image
            """
        if self.img_light_pth.count("_side_") > 0:
            self.img_light = self.img_light[:2900, :1000, :]
        else:
            self.img_light = self.img_light[:1700, :]
            # cv2.namedWindow("approxPolyDP", cv2.WINDOW_NORMAL)
            # cv.imshow("approxPolyDP", self.img_light)
            # cv.waitKey(0)
        """
        clean pic
        """
        img_gray = cv2.cvtColor(self.img_light, cv2.COLOR_BGR2GRAY)
        row, column = img_gray.shape
        left = 0
        right = 0
        for i in range(row):
            if np.sum(img_gray[i, :]) != 0:
                left = i
                break
        for i in range(row - 1, -1, -1):
            if np.sum(img_gray[i, :]) != 0:
                right = i + 1
                break
        img_gray = img_gray[left:right, :]
        self.img_light = self.img_light[left:right, :, :]
        row, column = img_gray.shape
        for i in range(column):
            if np.sum(img_gray[:, i]) != 0:
                left = i
                break
        for i in range(column - 1, -1, -1):
            if np.sum(img_gray[:, i]) != 0:
                right = i + 1
                break
        self.img_light = self.img_light[:, left:right, :]
        print(self.img_light.shape)

    def clean_height(self, nan_num=np.nan):
        """
        Auth: WZW
        根据nan，提取周围部分，并将剩余的nan数据转为nan_num
        UpArea: [ :1700, : ]
        SideArea: [ :2900, :1000]
        PreVersion: clean_height in V1
        Args:
            image: np.array
            nan_num:
        Returns:
            compressed data
        """
        if self.npy_height_pth.count("_side_") > 0:
            image = self.npy_height[:2900, :1000]
        else:
            image = self.npy_height[:1700, :]
        left = 0
        right = 0
        row, column = image.shape
        for i in range(row):
            if len(image[i, np.isnan(image[i, :])]) != column:
                left = i
                break
        for i in range(row - 1, -1, -1):
            if len(image[i, np.isnan(image[i, :])]) != column:
                right = i + 1
                break
        image = image[left:right, :]
        row, column = image.shape
        left = 0
        right = 0
        for i in range(column):
            if len(image[np.isnan(image[:, i]), i]) != row:
                left = i
                break
        for i in range(column - 1, -1, -1):
            if len(image[np.isnan(image[:, i]), i]) != row:
                right = i + 1
                break
        image = image[:, left:right]
        print(image.shape)
        if not np.isnan(nan_num):
            print("数据压缩，修改nan为：" + str(nan_num))
            nan_index = np.isnan(image)
            image[nan_index] = nan_num
        self.npy_height = image

    def npy_normalization_16_8(self, nan_num=np.nan):
        """
        Auth: WZW
        PreVersion: outputNorm in V1
        Args:
            npy_data: input numpy data
            nan_num: nan is changed to a number
            imgShow: show 16 bit image and 8 bit image or not
            npySave: save normed npy file or not
            npy_file: numpy file name, used to generate saved file name
            targetFolder: path to save npy file
        Returns:
            16 bit image,
            8 bit image,
            max num in point cloud
            min num in point cloud
        """
        z = self.npy_height.copy()
        x = z.copy()
        nan_index = np.isnan(z)

        # 找出不是nan数值中的最大最小值
        img_max = (x[~nan_index]).max()
        img_min = (x[~nan_index]).min()

        # 归一化产生16位灰度图
        img_16bit = (z - img_min) * 65535.0 / (img_max - img_min)
        # print(len(nan_index))

        # 设置nan
        if not np.isnan(nan_num):
            print("修改nan为：" + str(nan_num))
            img_16bit[nan_index] = nan_num
        # print(img_new_min)
        # print(img_new_max)
        # 产生对应的8位灰度图
        img_gray8 = img_16bit / 256
        # print(img_16bit.shape)
        self.img_16bit = img_16bit
        self.img_gray8 = img_gray8
        self.img_max = img_max
        self.img_min = img_min

        ui8 = img_gray8.astype(np.uint8)
        self.img_rgb = cv.cvtColor(ui8, cv.COLOR_GRAY2BGR)

    def image_channels_add_8bit(self):
        """

        Args:
            img_rgb:
            img_height:
            img_show:
            img_save:
            target_folder:
            file_name:

        Returns:

        """
        a, b = self.img_light.shape[:2]
        self.img_rgb = cv2.resize(self.img_rgb, (b, a))
        meta_img = np.zeros(self.img_light.shape)
        img_height_gray = cv2.cvtColor(self.img_rgb, cv2.COLOR_RGB2GRAY)
        img_rgb_gray = cv2.cvtColor(self.img_light, cv2.COLOR_RGB2GRAY)
        meta_img[:, :, 0] = img_height_gray
        meta_img[:, :, 1] = img_rgb_gray
        # meta_img[:, :, 2] = meta_img[:, :, 2] * 50
        meta_img = meta_img.astype(np.uint8)
        meta_img = cv2.cvtColor(meta_img, cv2.COLOR_BGR2RGB)
        # print(os.path.join(self.target_folder, os.path.basename(self.img_light_pth)[:-4] + "_channelsBlend.png"))
        cv2.imwrite(os.path.join(self.target_folder, os.path.basename(self.img_light_pth)[:-4] + "_channelsBlend.png"),
                    meta_img)

    def generate_manager(self):
        # start = time.time()
        try:
            res_dict = {}
            self.clean_light()
            self.clean_height()
            self.npy_normalization_16_8()
            self.image_channels_add_8bit()
            res_dict['image_output'] = os.path.join(self.target_folder, os.path.basename(self.img_light_pth)[
                                                                        :-4] + "_channelsBlend.png")
        except Exception as e:
            print(e)
            res_dict['error'] = e
        finally:
            if self.task_type == '110':
                res_dict['task_type'] = self.task_type
                res_dict['task_id'] = self.task_id
                send_message(self.MQ_send, self.connection_send, json.dumps(res_dict, ensure_ascii=False))
        return res_dict


if __name__ == "__main__":
    # # start0 = time.time()
    # # png = "D:\\Programs\\data\\dataset\\20220701-0628-full-size\\2022-06-28\\3505277134ML085B085301110012200465_side_2022-06-28-09-05-21_meta.png"
    # # npy = "D:\\Programs\\data\\dataset\\20220701-0628-full-size\\2022-06-28\\3505277134ML085B085301110012200465_side_2022-06-28-09-05-21.npy"
    # png = "D:\\Programs\\data\\dataset\\20220701-0628-full-size\\2022-06-28\\3505277134ML085B085301110012200465_up_2022-06-28-09-04-52_meta.png"
    # npy = "D:\\Programs\\data\\dataset\\20220701-0628-full-size\\2022-06-28\\3505277134ML085B085301110012200465_up_2022-06-28-09-04-52.npy"
    # test = preprocessing(0, 0, 0, 0, npy, png, 1, targetFolder='./out')
    # test.generate_manager()
    # # end0 = time.time()
    # # print(str(end0 - start0))

    ######## generate 2022-06-28
    # rootDir = 'D:\\Programs\\data\\dataset\\20220701-0628-full-size\\2022-06-28'
    # 3505277134ML085B085301118062200862_up_2022-08-09-10-43-58.npy
    npy = "D:\\Programs\\data\\dataset\\20220808+20220706\\2022-08-09\\3505277134ML085B085301118062200862_up_2022-08-09-10-43-58.npy"
    png = "D:\\Programs\\data\\dataset\\20220808+20220706\\2022-08-09\\3505277134ML085B085301118062200862_up_2022-08-09-10-43-58_meta.png"
    test = Preprocessing(0, 0, 0, 0, npy, png, 1, "D:\\Programs\\data\\exp\\pre")
    test.generate_manager()
    # rootDir = 'D:\\Programs\\data\\dataset\\20220823\\data'
    # file = os.listdir(rootDir)
    # for i in file:
    #     if i[-3:] == 'npy':
    #         print(i)
    #         start = time.time()
    #         npy = os.path.join(rootDir, i)
    #         png = os.path.join(rootDir, i[:-4] + "_meta.png")
    #         print(npy)
    #         print(png)
    #         # test = Preprocessing(0, 0, 0, 0, npy, png, 1,
    #         #                      targetFolder='D:\\Programs\\data\\dataset\\20220701-0628-full-size\\generate')
    #         test = Preprocessing(0, 0, 0, 0, npy, png, 1,
    #                              targetFolder='D:\\Programs\\data\\dataset\\20220823\\preprocessing')
    #         test.generate_manager()
    #         end = time.time()
    #         print(str(end - start))
